package net.fltx.icon;

import com.github.katjahahn.parser.PhysicalLocation;
import com.github.katjahahn.parser.sections.rsrc.icon.GroupIconDir;
import com.github.katjahahn.parser.sections.rsrc.icon.GroupIconDirEntry;
import com.github.katjahahn.parser.sections.rsrc.icon.GroupIconResource;
import com.github.katjahahn.parser.sections.rsrc.icon.IconParser;
import scala.collection.Iterator;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.util.List;

public class IconUtil {

    /**
     * 从pe文件中提取图标
     *
     * @param peFile exe,dll等PE文件
     * @return ico文件字节
     */
    public static byte[] ExtractIcon(String peFile) throws IOException {
        FileInputStream fis = null;
        ByteArrayOutputStream buf=null;
        try {
            File pe = new File(peFile);
            List<GroupIconResource> girs = IconParser.extractGroupIcons(pe);
            GroupIconResource gir = girs.get(0);
            GroupIconDir gid = gir.groupIconDir();
            IconDirEx iconDirEx = new IconDirEx((char) gid.idReserved(), (char) gid.idType(), (char) gid.idCount());
            Iterator<GroupIconDirEntry> iterator = gid.idEntries().iterator();
            int dataSize = 0;
            while (iterator.hasNext()) {
                GroupIconDirEntry next = iterator.next();
                IconDirEntryEx dirEntryEx = new IconDirEntryEx(next.bWidth(), next.bHeight(), next.bColorCount(), next.bReserved(),
                        (char) next.wPlanes(), (char) next.wBitCount(), (int) next.dwBytesInRes());
                iconDirEx.idEntries.add(dirEntryEx);
                dataSize += next.dwBytesInRes();
            }
            buf = new ByteArrayOutputStream();
            buf.write(iconDirEx.toByteArray());
            PhysicalLocation plo = gir.nIDToLocations().get(1).get();
            fis = new FileInputStream(pe);
            byte[] data = new byte[dataSize];
            fis.skip(plo.from());
            fis.read(data);
            buf.write(data);
            return buf.toByteArray();
        } catch (Exception e) {
            //e.printStackTrace();
            System.out.println("无法获取图标:"+peFile);
        } finally {
            if (fis!=null) fis.close();
            if (buf!=null) buf.close();
        }
        return null;
    }

}
